# -*- coding: utf-8 -*-

"""
# @Time    : 2025/3/3 14:17
# @User  : Mabin
# @Description  :MongoDB数据库操作工具类（单例、连接池）
"""
from bson import ObjectId
from pymongo import MongoClient
from threading import Lock

from application.settings import MONGODB_DATABASES


class MongoDBManager:
    """
    MongoDB数据库管理类
    :author Mabin
    使用示例：
    test_model = MongoDBManager()
    test_result = test_model.db["information"].insert_one({
        "name": "测试",
        "desc": "🤣🤣"
    })
    """
    _instances = {}  # 存储不同URI的连接实例
    _lock = Lock()  # 线程安全锁

    def __new__(cls, connect_key="default"):
        """
        单例核心逻辑
        :author Mabin
        :param str connect_key:数据库连接标识
        """
        if connect_key not in cls._instances:
            # 不存在对应连接标识的实例
            with cls._lock:
                # 实例化
                instance = super().__new__(cls)

                # 初始化数据库连接
                instance._init_connection(connect_key=connect_key)

                # 存储
                cls._instances[connect_key] = instance

        # 返回实例
        return cls._instances[connect_key]

    def _init_connection(self, connect_key="default"):
        """
        初始化新连接(连接池)
        :author Mabin
        :param str connect_key:数据库连接标识
        :return:
        """
        # 获取数据库链接配置
        connect_config = MONGODB_DATABASES.get(connect_key, None)
        if not connect_config:
            raise Exception(f"创建MongoDB数据库连接时，未查询到数据库链接配置！{connect_key}")

        # 创建连接池（自动管理连接池）
        self.client = MongoClient(
            host=connect_config["host"],
            port=connect_config["port"],
            username=connect_config["user"],
            password=connect_config["password"],
            authSource=connect_config["auth_source"],
            directConnection=True,
            maxPoolSize=200,  # 最大连接数
            minPoolSize=10,  # 最小保持连接数
            connectTimeoutMS=3000,  # 连接超时(3秒)
            socketTimeoutMS=5000,  # 操作超时(5秒)
            serverSelectionTimeoutMS=3000,  # 服务器选择超时
            waitQueueTimeoutMS=2000  # 等待队列超时
        )
        self.db = self.client[connect_config["database"]]  # 切换至指定数据库


class MongoDBDataStream:
    """
    MongoDB 数据流迭代器类，读取时将 `None` 替换为空字符串 ""。
    """

    mongodb_manager = MongoDBManager()

    def __init__(self, collection: str, batch_size: int, sort_key: str, historical_cursor_position: str):
        """
        初始化数据流读取器

        :param collection: 目标集合名称
        :param batch_size: 每次从服务端拉取的文档数量
        :param sort_key: 用作排序的字段（通常为 "_id"）
        :param historical_cursor_position: 历史游标位置（上次同步的 `_id` 的字符串表示），为空表示全量读取
        """
        self.collection = collection
        self.batch_size = batch_size
        self.sort_key = sort_key

        # 只有当 historical_cursor_position 不为空时才创建游标过滤条件
        if historical_cursor_position:
            try:
                self.cursor_filter = {
                    sort_key: {"$gt": ObjectId(historical_cursor_position)}
                }
            except Exception as e:
                # 如果游标位置无效，则回退到全量读取模式
                self.cursor_filter = {}
        else:
            self.cursor_filter = {}

    @staticmethod
    def _replace_none_with_empty(obj):
        """
        递归将对象中所有值为 None 的位置替换为空字符串 ""。

        支持的结构：dict, list, tuple；其他类型原样返回。
        :param obj: 待处理对象
        :return: 处理后的对象（对于可变对象 dict/list 会返回新的对象或新的结构）
        """
        # 基本类型为 None -> 替换为空字符串
        if obj is None:
            return ""
        # dict：对每个键值递归处理
        if isinstance(obj, dict):
            return {k: MongoDBDataStream._replace_none_with_empty(v) for k, v in obj.items()}
        # list：递归处理每个元素
        if isinstance(obj, list):
            return [MongoDBDataStream._replace_none_with_empty(v) for v in obj]
        # tuple：返回新的 tuple
        if isinstance(obj, tuple):
            return tuple(MongoDBDataStream._replace_none_with_empty(v) for v in obj)
        # 其他类型原样返回（例如 ObjectId、int、str、datetime 等）
        return obj

    def stream(self, query: dict = None, fields: dict = None, max_limit: int = None):
        """
        按条件流式返回 MongoDB 文档（generator），将 None 值转换为 ""。

        :param query: 额外查询条件（字典形式）
        :param fields: 投影字段（字典），例如 {"name": 1, "_id": 0}
        :param max_limit: 最大返回文档数量（None 表示不限制）
        :yield: 逐条返回处理后的文档（所有 None 均替换为 ""，且 `_id` 转为字符串）
        """
        # 合并用户查询与增量游标条件
        final_filter = (query or {}) | self.cursor_filter

        cursor = (
            self.mongodb_manager.db[self.collection]
            .find(filter=final_filter, projection=fields)
            .sort(self.sort_key, 1)
            .batch_size(self.batch_size)
        )

        # 如果指定了最大数量，调用 limit
        if max_limit is not None:
            cursor = cursor.limit(max_limit)

        for raw_doc in cursor:
            _id_val = raw_doc.get('_id', None)
            processed = self._replace_none_with_empty(raw_doc)
            if _id_val is not None:
                processed['_id'] = str(_id_val)
            yield processed


if __name__ == '__main__':
    test_model = MongoDBDataStream(collection='raw_information_list_temp', batch_size=1000, sort_key='_id',
                                   historical_cursor_position='')
    for doc in test_model.stream(query={}, fields={}):
        print(doc)
